'''
Created on Nov 4, 2009

@author: mkiyer
'''

import sys
from bx.intervals.intersection import Interval, IntervalTree
import collections
import operator

def parse_raw_interval_file(fhd, category):
    for id, line in enumerate(fhd):
        fields = line.strip().split('\t')
        chrom = fields[0]
        start = int(fields[1])
        end = int(fields[2])
        yield Interval(start, end, value=('%s.%d' % (category, id+1), category), chrom=chrom)

def parse_bed_file(fhd, category):
    for line in fhd:
        fields = line.strip().split('\t')
        chrom = fields[0]
        start = int(fields[1])
        end = int(fields[2])
        name = fields[3]
        strand = fields[5]
        yield Interval(start, end, value=(name, category), chrom=chrom, strand=strand)

class Gene:    
    def __init__(self, interval):
        self.name = interval.value[0]
        self.chrom = interval.chrom
        self.start = interval.start
        self.end = interval.end
        self.exons = []
        
if __name__ == '__main__':    
    genes = {}
    totalgenes = 0
    trees = collections.defaultdict(lambda: IntervalTree())
    for interval in parse_raw_interval_file(open(sys.argv[1]), 'LINCRNA'):
        totalgenes += 1
        trees[interval.chrom].insert_interval(interval)
        genes[interval.value[0]] = Gene(interval)

    nohits = 0
    multiplehits = 0
    totalexons = 0
    
    for interval in parse_raw_interval_file(open(sys.argv[2]), 'lincrna_exon'):
        totalexons += 1
        hits = trees[interval.chrom].find(interval.start, interval.end)
        if len(hits) == 0:
            print 'no hits!'
            nohits += 1
            continue
        if len(hits) > 1:
            print 'wack!', interval.start, interval.end
            multiplehits += 1
            continue        
        name = hits[0].value[0]
        gene = genes[name]
        gene.exons.append((interval.start, interval.end))
        
    print '#nohits', nohits
    print '#multiplehits', multiplehits
    print '#totalexons', totalexons
    print '#totalgenes', totalgenes
    print '#exons per gene', float(totalexons)/totalgenes
    
    #sys.exit(0)
    source = 'lincRNA-literature'
    
    for gname in sorted(genes.iterkeys()):
        g = genes[gname]
        #chr1    Cufflinks       transcript      5180    6595    1000    -       .       gene_id "3064YAAXX_2.26"; transcript_id "3
        #064YAAXX_2.26.0"; RPKM "1.4581570540"; frac "1.000000"; conf_lo "0.000000"; conf_hi "0.000000"; cov "0.875886";
        #chr1    Cufflinks       exon    5180    6173    1000    -       .       gene_id "3064YAAXX_2.26"; transcript_id "3064YAAXX
        #_2.26.0"; exon_number "1"; RPKM "1.4581570540"; frac "1.000000"; conf_lo "0.000000"; conf_hi "0.000000"; cov "0.875886";
        #chr1    Cufflinks       exon    6470    6595    1000    -       .       gene_id "3064YAAXX_2.26"; transcript_id "3064YAAXX
        #_2.26.0"; exon_number "2"; RPKM "1.4581570540"; frac "1.000000"; conf_lo "0.000000"; conf_hi "0.000000"; cov "0.875886";
        gene_id = g.name
        transcript_id = gene_id + '.0'
        
        if len(g.exons) == 0:
            gtf_fields = [g.chrom,
                          source,
                          'exon',
                          str(g.start),
                          str(g.end),
                          '0.00',
                          '+',
                          '.',
                          'gene_id "%s"; transcript_id "%s";' % (gene_id, transcript_id)]
            print '\t'.join(gtf_fields)
        else:
            gtf_fields = [g.chrom,
                          source,
                          'transcript',
                          str(g.start),
                          str(g.end),
                          '0.00',
                          '+',
                          '.',
                          'gene_id "%s"; transcript_id "%s";' % (gene_id, transcript_id)]
            print '\t'.join(gtf_fields)

            sorted_exons = sorted(g.exons, key=operator.itemgetter(0))
            for e in sorted_exons:
                gtf_fields = [g.chrom,
                              source,
                              'exon',
                              str(e[0]),
                              str(e[1]),
                              '0.00',
                              '+',
                              '.',
                              'gene_id "%s"; transcript_id "%s";' % (gene_id, transcript_id)]
                print '\t'.join(gtf_fields)
        
        
